#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Data loading and processing module for item-based collaborative filtering recommendation system.

Created on: 2025-04-20
Author: Nianqing Liu
"""

import os
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split


class DataLoader:
    """
    Class for loading and preprocessing the MovieLens dataset.
    """

    def __init__(self, data_path="data/ml-100k"):
        """
        Initialize the DataLoader with path to the MovieLens data.

        Parameters:
        -----------
        data_path : str
            Path to the directory containing the MovieLens dataset files.
        """
        self.data_path = data_path
        self.ratings_df = None
        self.movies_df = None
        self.train_data = None
        self.test_data = None
        self.n_users = None
        self.n_items = None
        self.ratings_matrix = None
        self.item_user_matrix = None  # Added for item-based CF
        self.train_matrix = None
        self.test_matrix = None

    def load_data(self):
        """
        Load the MovieLens dataset from files.

        Returns:
        --------
        tuple
            A tuple containing (ratings_df, movies_df)
        """
        # Check if data directory exists
        if not os.path.exists(self.data_path):
            raise FileNotFoundError(f"Data directory {self.data_path} not found.")

        # Load ratings data
        ratings_file = os.path.join(self.data_path, "u.data")
        if not os.path.exists(ratings_file):
            raise FileNotFoundError(f"Ratings file {ratings_file} not found.")

        column_names = ["user_id", "item_id", "rating", "timestamp"]
        self.ratings_df = pd.read_csv(ratings_file, sep="\t", names=column_names)

        # Load movie data
        movies_file = os.path.join(self.data_path, "u.item")
        if not os.path.exists(movies_file):
            raise FileNotFoundError(f"Movies file {movies_file} not found.")

        movies_column_names = [
            "movie_id",
            "movie_title",
            "release_date",
            "video_release_date",
            "IMDb_URL",
            "unknown",
            "Action",
            "Adventure",
            "Animation",
            "Children",
            "Comedy",
            "Crime",
            "Documentary",
            "Drama",
            "Fantasy",
            "Film-Noir",
            "Horror",
            "Musical",
            "Mystery",
            "Romance",
            "Sci-Fi",
            "Thriller",
            "War",
            "Western",
        ]
        self.movies_df = pd.read_csv(
            movies_file, sep="|", names=movies_column_names, encoding="latin-1"
        )

        # Get number of users and items
        self.n_users = self.ratings_df["user_id"].nunique()
        self.n_items = self.ratings_df["item_id"].nunique()

        return self.ratings_df, self.movies_df

    def create_matrix(self):
        """
        Create user-item rating matrix from the ratings dataframe.

        Returns:
        --------
        numpy.ndarray
            User-item rating matrix.
        """
        if self.ratings_df is None:
            self.load_data()

        # Get the maximum user_id and item_id instead of just counting unique values
        max_user_id = self.ratings_df["user_id"].max()
        max_item_id = self.ratings_df["item_id"].max()

        # Update the dimensions
        self.n_users = max_user_id
        self.n_items = max_item_id

        # Initialize matrix with zeros (using maximum IDs)
        self.ratings_matrix = np.zeros((self.n_users, self.n_items))

        # Fill matrix with ratings
        for row in self.ratings_df.itertuples():
            # Adjust indices to 0-based
            self.ratings_matrix[row.user_id - 1, row.item_id - 1] = row.rating

        return self.ratings_matrix

    def create_item_user_matrix(self):
        """
        Create item-user rating matrix (transpose of user-item matrix).
        Useful for item-based collaborative filtering.

        Returns:
        --------
        numpy.ndarray
            Item-user rating matrix.
        """
        if self.ratings_matrix is None:
            self.create_matrix()

        # Transpose user-item matrix to get item-user matrix
        self.item_user_matrix = self.ratings_matrix.T

        return self.item_user_matrix

    def split_data(self, test_size=0.2, random_state=42):
        """
        Split the ratings data into training and testing sets.

        Parameters:
        -----------
        test_size : float
            Proportion of the dataset to include in the test split (default 0.2).
        random_state : int
            Random seed for reproducibility (default 42).

        Returns:
        --------
        tuple
            A tuple containing (train_data, test_data, train_matrix, test_matrix)
        """
        if self.ratings_df is None:
            self.load_data()

        # Make sure dimensions are set correctly
        if self.n_users is None or self.n_items is None:
            max_user_id = self.ratings_df["user_id"].max()
            max_item_id = self.ratings_df["item_id"].max()
            self.n_users = max_user_id
            self.n_items = max_item_id

        # Split ratings dataframe
        self.train_data, self.test_data = train_test_split(
            self.ratings_df, test_size=test_size, random_state=random_state
        )

        # Create training matrix
        self.train_matrix = np.zeros((self.n_users, self.n_items))
        for row in self.train_data.itertuples():
            self.train_matrix[row.user_id - 1, row.item_id - 1] = row.rating

        # Create test matrix (for evaluation)
        self.test_matrix = np.zeros((self.n_users, self.n_items))
        for row in self.test_data.itertuples():
            self.test_matrix[row.user_id - 1, row.item_id - 1] = row.rating

        return self.train_data, self.test_data, self.train_matrix, self.test_matrix

    def get_user_ratings(self, user_id):
        """
        Get all ratings for a specific user.

        Parameters:
        -----------
        user_id : int
            ID of the user to get ratings for.

        Returns:
        --------
        pandas.DataFrame
            DataFrame with the user's ratings joined with movie information.
        """
        if self.ratings_df is None or self.movies_df is None:
            self.load_data()

        # Get user ratings
        user_ratings = self.ratings_df[self.ratings_df["user_id"] == user_id]

        # Join with movie information
        user_ratings_with_info = user_ratings.merge(
            self.movies_df[["movie_id", "movie_title"]],
            left_on="item_id",
            right_on="movie_id",
        )

        return user_ratings_with_info

    def get_item_ratings(self, item_id):
        """
        Get all ratings for a specific item.

        Parameters:
        -----------
        item_id : int
            ID of the item to get ratings for.

        Returns:
        --------
        pandas.DataFrame
            DataFrame with the item's ratings joined with user information.
        """
        if self.ratings_df is None:
            self.load_data()

        # Get item ratings
        item_ratings = self.ratings_df[self.ratings_df["item_id"] == item_id]

        # If movie information is available, include it
        if self.movies_df is not None:
            movie_info = self.movies_df[self.movies_df["movie_id"] == item_id]
            if not movie_info.empty:
                movie_title = movie_info.iloc[0]["movie_title"]
                item_ratings = item_ratings.assign(movie_title=movie_title)

        return item_ratings

    def get_dataset_stats(self):
        """
        Get basic statistics about the dataset.

        Returns:
        --------
        dict
            Dictionary containing dataset statistics.
        """
        if self.ratings_df is None or self.movies_df is None:
            self.load_data()

        # Calculate density
        density = len(self.ratings_df) / (self.n_users * self.n_items)

        stats = {
            "n_users": self.n_users,
            "n_items": self.n_items,
            "n_ratings": len(self.ratings_df),
            "rating_density": density,
            "rating_min": self.ratings_df["rating"].min(),
            "rating_max": self.ratings_df["rating"].max(),
            "rating_mean": self.ratings_df["rating"].mean(),
            "rating_median": self.ratings_df["rating"].median(),
        }

        return stats

    def get_item_stats(self):
        """
        Calculate statistics for items (movies).
        Useful for item-based CF analysis.

        Returns:
        --------
        pandas.DataFrame
            DataFrame with item statistics (rating counts, average ratings).
        """
        if self.ratings_df is None:
            self.load_data()

        # Group by item_id and calculate stats
        item_stats = (
            self.ratings_df.groupby("item_id")
            .agg(
                rating_count=("rating", "count"),
                rating_mean=("rating", "mean"),
                rating_std=("rating", "std"),
            )
            .reset_index()
        )

        # Join with movie information if available
        if self.movies_df is not None:
            item_stats = item_stats.merge(
                self.movies_df[["movie_id", "movie_title"]],
                left_on="item_id",
                right_on="movie_id",
            )

        return item_stats
